#include "hero_integrated_vision/traditional_detector.hpp"
#include <cmath>
#include <algorithm>
#include <sstream>
namespace rm_aim
{
  TraditionalDetector::TraditionalDetector(
      const rclcpp::Logger &logger,
      const ArmorParams &armor_params,
      const LightParams &light_params,
      double conf_threshold,
      int bin_thres,
      int detect_color, 
      const std::vector<std::string> &ignore_classes)
      : BaseDetector(logger, conf_threshold),
        armor_params_(armor_params),
        light_params_(light_params),
        binary_thres_(bin_thres),
        detect_color_(detect_color),
        ignore_classes_(ignore_classes)
  {
  }

  bool TraditionalDetector::init(const std::string &model_path, const std::string &device)
  {
    try
    {
      // 初始化数字分类器（参数：模型路径、空标签路径、置信度阈值、日志器）
      classifier_ = std::make_unique<NumberClassifier>(
          model_path,
          "",
          conf_threshold_,
          logger_,
          ignore_classes_);
      // 编译模型（指定运行设备：CPU/CUDA）
      classifier_->compileModel(device);
      // 标记初始化完成
      is_initialized_ = true;
      RCLCPP_INFO(logger_, "TraditionalDetector initialized successfully (device: %s)", device.c_str());
      return true;
    }
    catch (const std::exception &e)
    {
      // 捕获初始化异常（如模型路径错误、设备不支持）
      RCLCPP_ERROR(logger_, "TraditionalDetector init failed: %s", e.what());
      is_initialized_ = false;
      return false;
    }
  }
  bool TraditionalDetector::init(const std::string &model_path, const std::string &label_path, const std::string &device)
  {
    try
    {
      // 初始化数字分类器（参数：模型路径、空标签路径、置信度阈值、日志器）
      classifier_ = std::make_unique<NumberClassifier>(
          model_path,
          label_path,
          conf_threshold_,
          logger_,
          ignore_classes_);
      // 编译模型（指定运行设备：CPU/CUDA）
      classifier_->compileModel(device);
      // 标记初始化完成
      is_initialized_ = true;
      RCLCPP_INFO(logger_, "TraditionalDetector initialized successfully (device: %s)", device.c_str());
      return true;
    }
    catch (const std::exception &e)
    {
      // 捕获初始化异常（如模型路径错误、设备不支持）
      RCLCPP_ERROR(logger_, "TraditionalDetector init failed: %s", e.what());
      is_initialized_ = false;
      return false;
    }
  }

  std::vector<Armor> TraditionalDetector::detect(const cv::Mat &rgb_image)
  {
    // 检查检测器就绪状态（未初始化则返回空列表）
    if (!isReady())
    {
      RCLCPP_WARN(logger_, "TraditionalDetector not ready (init incomplete or classifier null)");
      return {};
    }

    // 步骤1：图像预处理（RGB→灰度→二值化）
    cv::Mat binary_img = preprocessImage(rgb_image);

    // 步骤2：灯条检测（从二值图中筛选有效灯条）
    lights_ = findLights(rgb_image, binary_img);

    // 步骤3：灯条匹配（生成有效装甲板）
    armors_ = matchLights(lights_);

    // 步骤4：数字分类（若有装甲板则提取数字并分类）
    if (!armors_.empty() && classifier_ != nullptr)
    {
      classifier_->extractNumbers(rgb_image, armors_); // 提取装甲板数字ROI
      classifier_->classify(armors_);                  // 分类数字并填充结果
    }

    return armors_;
  }

  bool TraditionalDetector::isReady() const
  {
    // 需同时满足：初始化完成 + 数字分类器有效
    return is_initialized_ && (classifier_ != nullptr);
  }

  cv::Mat TraditionalDetector::preprocessImage(const cv::Mat &rgb_image)
  {
    // 1. RGB图像转灰度图（减少通道数，简化计算）
    cv::Mat gray_img;
    cv::cvtColor(rgb_image, gray_img, cv::COLOR_RGB2GRAY);

    // 2. 灰度图二值化（突出灯条区域，抑制背景）
    cv::Mat binary_img;
    // THRESH_BINARY：灰度值>binary_thres_设为255（白），否则设为0（黑）
    cv::threshold(gray_img, binary_img, binary_thres_, 255, cv::THRESH_BINARY);

    return binary_img;
  }

  std::vector<Light> TraditionalDetector::findLights(const cv::Mat &rgb_image, const cv::Mat &binary_image)
  {
    std::vector<Light> valid_lights;              // 存储最终有效灯条
    std::vector<std::vector<cv::Point>> contours; // 存储轮廓
    std::vector<cv::Vec4i> hierarchy;             // 存储轮廓层级（仅保留外层轮廓）

    // 1. 提取二值图中的外层轮廓（RETR_EXTERNAL：只找最外层轮廓）
    cv::findContours(
        binary_image,
        contours,
        hierarchy,
        cv::RETR_EXTERNAL,
        cv::CHAIN_APPROX_SIMPLE // 简化轮廓（减少点数量）
    );

    // 2. 遍历每个轮廓，筛选有效灯条
    for (const auto &contour : contours)
    {
      // 轮廓点数<5时，无法拟合稳定的旋转矩形，直接跳过
      if (contour.size() < 5)
        continue;

      // 3. 拟合轮廓的最小面积矩形（灯条的几何形态）
      cv::RotatedRect min_area_rect = cv::minAreaRect(contour);
      // 用旋转矩形初始化灯条对象（自动计算长度、宽度、倾斜角等参数）
      Light candidate_light(min_area_rect);

      // 4. 判断灯条是否有效（宽高比、倾斜角符合要求）
      if (isValidLight(candidate_light))
      {
        // 5. 检查灯条包围盒是否超出图像边界（避免后续ROI访问越界）
        cv::Rect light_rect = candidate_light.boundingRect();
        if (light_rect.x >= 0 && light_rect.y >= 0 &&
            light_rect.x + light_rect.width <= rgb_image.cols &&
            light_rect.y + light_rect.height <= rgb_image.rows)
        {
          // 6. 统计灯条区域的红/蓝色像素总和（判断阵营颜色）
          int sum_blue = 0, sum_red = 0;
          cv::Mat light_roi = rgb_image(light_rect); // 灯条ROI区域

          // 遍历ROI像素，仅统计轮廓内的像素（避免背景干扰）
          for (int row = 0; row < light_roi.rows; ++row)
          {
            for (int col = 0; col < light_roi.cols; ++col)
            {
              // 转换像素坐标到原图坐标系，判断是否在轮廓内
              cv::Point2f original_point(col + light_rect.x, row + light_rect.y);
              if (cv::pointPolygonTest(contour, original_point, false) >= 0)
              {
                // RGB图像通道顺序：0=R，1=G，2=B
                sum_red += light_roi.at<cv::Vec3b>(row, col)[0];  // 红色通道总和
                sum_blue += light_roi.at<cv::Vec3b>(row, col)[2]; // 蓝色通道总和
              }
            }
          }

          // 7. 根据红/蓝色总和判断灯条阵营（红>蓝为RED_CAMP，反之BLUE_CAMP）
          candidate_light.camp_color = (sum_red > sum_blue) ? RED_CAMP : BLUE_CAMP;
          // 将有效灯条加入列表
          valid_lights.emplace_back(candidate_light);
        }
      }
    }

    return valid_lights;
  }

  bool TraditionalDetector::isValidLight(const Light &candidate_light)
  {
    // 1. 计算灯条宽高比（width/length，横向尺寸/纵向尺寸）
    double aspect_ratio = candidate_light.width_pixel / candidate_light.length_pixel;
    // 宽高比需在[min_aspect_ratio, max_aspect_ratio]范围内
    bool is_aspect_ratio_valid = (aspect_ratio > light_params_.min_aspect_ratio) &&
                                 (aspect_ratio < light_params_.max_aspect_ratio);

    // 2. 灯条倾斜角需小于最大竖直倾斜角
    bool is_tilt_angle_valid = (candidate_light.tilt_angle_deg < light_params_.max_vertical_tilt_angle);

    // 综合判断：宽高比和倾斜角均有效才是有效灯条
    return is_aspect_ratio_valid && is_tilt_angle_valid;
  }

  std::vector<Armor> TraditionalDetector::matchLights(const std::vector<Light> &lights)
  {
    std::vector<Armor> valid_armors; // 存储最终有效装甲板

    // 双重循环遍历所有灯条对（避免重复匹配）
    for (size_t i = 0; i < lights.size(); ++i)
    {
      for (size_t j = i + 1; j < lights.size(); ++j)
      {
        const Light &light1 = lights[i];
        const Light &light2 = lights[j];

        // 1. 过滤：灯条颜色需与目标检测颜色一致
        if (light1.camp_color != detect_color_ || light2.camp_color != detect_color_)
          continue;

        // 2. 过滤：灯条对包围盒内不能包含其他灯条（避免多灯条误判）
        if (isLightPairContainOtherLight(light1, light2, lights))
          continue;

        // 3. 判断灯条对是否构成有效装甲板，并获取装甲板类型
        ArmorType armor_type = judgeArmorType(light1, light2);
        if (armor_type != ArmorType::INVALID)
        {
          // 4. 构造装甲板对象（自动按x坐标排序左右灯条）
          Armor valid_armor(light1, light2);
          valid_armor.type = armor_type; // 填充装甲板类型
          valid_armors.emplace_back(valid_armor);
        }
      }
    }

    return valid_armors;
  }

  bool TraditionalDetector::isLightPairContainOtherLight(const Light &light1, const Light &light2, const std::vector<Light> &lights)
  {
    // 1. 获取灯条对的四个关键顶点（两个灯条的上下顶点）
    std::vector<cv::Point2f> key_points = {
        light1.top_point, light1.bottom_point,
        light2.top_point, light2.bottom_point};

    // 2. 计算四个顶点的包围矩形（覆盖灯条对的最小矩形）
    cv::Rect bounding_rect = cv::boundingRect(key_points);

    // 3. 检查所有其他灯条是否在包围矩形内
    for (const auto &test_light : lights)
    {
      // 跳过当前灯条对（只检查其他灯条）
      if (test_light.center == light1.center || test_light.center == light2.center)
        continue;

      // 若其他灯条的中心/上顶点/下顶点在包围矩形内，说明包含其他灯条
      if (bounding_rect.contains(test_light.center) ||
          bounding_rect.contains(test_light.top_point) ||
          bounding_rect.contains(test_light.bottom_point))
      {
        return true;
      }
    }

    return false;
  }

  ArmorType TraditionalDetector::judgeArmorType(const Light &light1, const Light &light2)
  {
    // 1. 检查灯条长度比（短灯条长度/长灯条长度）
    double light_length_ratio = (light1.length_pixel < light2.length_pixel) ? (light1.length_pixel / light2.length_pixel) : (light2.length_pixel / light1.length_pixel);
    bool is_light_length_valid = (light_length_ratio > armor_params_.min_light_length_ratio);

    // 2. 计算灯条中心距与平均长度的比值（归一化，避免尺寸影响）
    double avg_light_length = (light1.length_pixel + light2.length_pixel) / 2.0;
    double center_distance = cv::norm(light1.center - light2.center); // 灯条中心欧氏距离
    double normalized_center_dist = center_distance / avg_light_length;

    // 3. 检查中心距是否在小型/大型装甲板的范围内
    bool is_small_armor_dist_valid = (normalized_center_dist >= armor_params_.min_small_armor_center_dist) &&
                                     (normalized_center_dist < armor_params_.max_small_armor_center_dist);
    bool is_large_armor_dist_valid = (normalized_center_dist >= armor_params_.min_large_armor_center_dist) &&
                                     (normalized_center_dist < armor_params_.max_large_armor_center_dist);
    bool is_center_dist_valid = is_small_armor_dist_valid || is_large_armor_dist_valid;

    // 4. 检查灯条中心连线的水平倾斜角（与水平方向的夹角）
    cv::Point2f center_diff = light1.center - light2.center; // 中心差值（x方向为水平，y为竖直）
    double angle_deg = std::abs(std::atan2(center_diff.y, center_diff.x)) * 180 / CV_PI;
    double horizontal_tilt_angle = (angle_deg > 90) ? (180 - angle_deg) : angle_deg;
    bool is_tilt_angle_valid = (horizontal_tilt_angle < armor_params_.max_horizontal_tilt_angle);

    // 5. 综合判断：所有条件满足则为有效装甲板，否则为无效
    if (!(is_light_length_valid && is_center_dist_valid && is_tilt_angle_valid))
    {
      return ArmorType::INVALID;
    }

    // 6. 判断装甲板类型（中心距在小型范围则为SMALL，否则为LARGE）
    return (is_small_armor_dist_valid) ? ArmorType::SMALL : ArmorType::LARGE;
  }

  void TraditionalDetector::drawArmorDetectionResults(cv::Mat &img)
  {
    for (const auto &armor : armors_)
    {
      // 根据灯条颜色绘制装甲板，红色灯条用青色（255,255,0），蓝色灯条用紫红色（255,0,255）
      auto line_color = armor.left_light.camp_color == RED_CAMP ? cv::Scalar(255, 255, 0) : cv::Scalar(255, 0, 255);
      // 绘制装甲板框
      tools::draw_points(img, armor.apex, line_color, 1);
      // 绘制文本，显示装甲板的颜色、名称和置信度
      tools::draw_text(img, armor.classification_result, armor.left_light.top_point, cv::Scalar(0, 255, 255), 0.8, 2);
    }
  }

} // namespace rm_aim